#ifndef __CUHB_RNG_CUH__
#define __CUHB_RNG_CUH__

/*
  Copyright (c) 2010 Davide Rossetti (davide.rossetti@roma1.infn.it)

  This file is part of CuHB (CUDA Heisenberg) package.
  
  CuHB is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  CuHB is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  
  You should have received a copy of the GNU General Public License
  along with CuHB.  If not, see <http://www.gnu.org/licenses/>.

*/

#include <stdlib.h> // for RAND_MAX
#include <math_constants.h> // for CUDART_PI_F

#include "rng.hpp"

// #define RANDOM ((ira[ip++]=ira[ip1++]+ira[ip2++])^ira[ip3++])
// unsigned char ip, ip1, ip2, ip3;
// ip=255;
// ip1=ip-24;
// ip2=ip-55;
// ip3=ip-61;

// WATCH OUT!!! 256 here can't be changed easily
// it is linked to 8bit char type of ip below, and it's auto wrap-around feature
#define RNG_N 256

#define make_rng_int2(a,b)     make_uint2(a,b)
#define make_rng_int4(a,b,c,d) make_uint4(a,b,c,d)

typedef struct rng_state {
    uchar4 ip;
    rng_int_t ira[RNG_N];
} rng_state;

//#define RNG_INT_MAX RAND_MAX
//#define RNG_INT_MAX  (~(rng_int_t)0)
#define RNG_INT_MAX  (rng_int_t)0XFFFFFFFF
#define RNG_INT_MIN  (rng_int_t)0X0


//-----------------------------------------------------------------------------
// Uniform [0,1]
//-----------------------------------------------------------------------------

__device__ rng_int_t rng_uniform_int(rng_state* states)
{
    const int tid = threadIdx.x;
    const int bdim = blockDim.x;
    int i;
    rng_state* state = &states[blockIdx.x];
    rng_int_t rng = 0;
    for(i=0; i<bdim; ++i) {
        if(i==tid) {
            uchar4 ip = state->ip;
            rng_int_t tmp;

            tmp = state->ira[ip.y]+state->ira[ip.z];
            rng = tmp ^ state->ira[ip.w];
            state->ira[ip.x] = tmp;
            ip.x++; ip.y++; ip.z++; ip.w++;
            state->ip = ip;
        }
        __syncthreads();
    }
    return rng;
}

//-----------------------------------------------------------------------------

__device__ rng_int_t rng_uniform_shm_int(rng_state* states)
{
    //const int bid = blockIdx.x;
    const int tid = threadIdx.x;
    const int bdim = blockDim.x;
    //const int idx = bid*bdim+tid;
    int i;
    rng_state* state = &states[blockIdx.x];
    rng_int_t rng = 0;
    __shared__ int shared_ira[RNG_N];
    __shared__ uchar4 shared_ip;
    int bs;

    // copy state to shared state
    if(tid == 0) {
        shared_ip = state->ip;
    }
    __syncthreads();
    // int lim = RNG_N/bdim;
    // lim = lim>0?lim:1;
    //for(bs=0; bs<lim; ++bs) {
    // ....
    //}
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            shared_ira[ptr] = state->ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    for(i=0; i<bdim; ++i) {
        if(i==tid) {
            uchar4 ip = shared_ip;
            rng_int_t tmp = shared_ira[ip.y]+shared_ira[ip.z];
            shared_ira[ip.x] = tmp;
            rng = tmp ^ shared_ira[ip.w];
            ip.x++;
            ip.y++;
            ip.z++;
            ip.w++;
            shared_ip = ip;
        }
        __syncthreads();
    }

    // copy back shared state
    if(tid == 0) {
        state->ip = shared_ip;
    }
    __syncthreads();
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            state->ira[ptr] = shared_ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    return rng;
}

//-----------------------------------------------------------------------------
// parallel calc of BB threads

// assert BB % bdim == 0!!!
// assert BB < 24
#define BB 16

__device__ rng_int_t rng_uniform_shm16_int(rng_state* states)
{
    const int tid = threadIdx.x;
    const int bdim = blockDim.x;
    int i;
    rng_state* state = &states[blockIdx.x];
    rng_int_t rng = 0;
    __shared__ int shared_ira[RNG_N];
    __shared__ uchar4 shared_ip;
    int bs;

    // copy state to shared state
    if(tid == 0) {
        shared_ip = state->ip;
    }
    __syncthreads();

    // loop only if RNG_N > bdim(==BLOCK_SIZE)
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            shared_ira[ptr] = state->ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    for(i=0; i<bdim; i+=BB) {
        int off = tid - i;
        uchar4 ip = shared_ip;
        if(0 <= off && off < BB) {
            rng_int_t tmp = shared_ira[ip.y+off]+shared_ira[ip.z+off];
            shared_ira[ip.x+off] = tmp;
            rng = tmp ^ shared_ira[ip.w+off];
            //ip = make_uchar4(ip.x+off,ip.y+off,ip.z+off,ip.w+off);
        }
        __syncthreads();
        if(off == 0) {
            shared_ip = make_uchar4(ip.x+BB,ip.y+BB,ip.z+BB,ip.w+BB);
        }
        __syncthreads();
    }

    // copy back shared state
    if(tid == 0) {
        state->ip = shared_ip;
    }
    __syncthreads();
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            state->ira[ptr] = shared_ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    return rng;
}

//-----------------------------------------------------------------------------

__device__ rng_int2_t rng_uniform_shm16_int2(rng_state* states)
{
    const int tid = (threadIdx.x * blockDim.y + threadIdx.y) * blockDim.z + threadIdx.z; 
    const int bdim = blockDim.x * blockDim.y * blockDim.z;
    int i;
    rng_state* state = &states[blockIdx.x];
    rng_int2_t ret;
    __shared__ int shared_ira[RNG_N];
    __shared__ uchar4 shared_ip;
    int bs;

    // copy state to shared state
    if(tid == 0) {
        shared_ip = state->ip;
    }
    __syncthreads();

    // loop only if RNG_N > bdim(==BLOCK_SIZE)
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            shared_ira[ptr] = state->ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    ret = make_rng_int2(0,0);
    for(int nw=0; nw<2; ++nw) {
        rng_int_t rng = 0;
        for(i=0; i<bdim; i+=BB) {
            int off = tid - i;
            uchar4 ip = shared_ip;
            if(0 <= off && off < BB) {
                rng_int_t tmp = shared_ira[ip.y+off]+shared_ira[ip.z+off];
                shared_ira[ip.x+off] = tmp;
                rng = tmp ^ shared_ira[ip.w+off];
                ret.x = ret.y;
                ret.y = rng;
                //ip = make_uchar4(ip.x+off,ip.y+off,ip.z+off,ip.w+off);
            }
            __syncthreads();

            if(off == 0)
                shared_ip = make_uchar4(ip.x+BB,ip.y+BB,ip.z+BB,ip.w+BB);
            __syncthreads();
        }
    }

    // copy back shared state
    if(tid == 0) {
        state->ip = shared_ip;
    }
    __syncthreads();
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            state->ira[ptr] = shared_ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    //return make_uint2(rng1,rng2);
    return ret;
}

//-----------------------------------------------------------------------------

__device__ rng_int4_t rng_uniform_shm16_int4(rng_state* states)
{
    //const int bid = blockIdx.x;
    const int tid = threadIdx.x;
    const int bdim = blockDim.x;
    //const int idx = bid*bdim+tid;
    int i;
    rng_state* state = &states[blockIdx.x];
    //uint rng1 = 0, rng2 = 0;
    rng_int4_t ret;
    __shared__ int shared_ira[RNG_N];
    __shared__ uchar4 shared_ip;
    int bs;

    // copy state to shared state
    if(tid == 0) {
        shared_ip = state->ip;
    }
    __syncthreads();

    // loop only if RNG_N > bdim(==BLOCK_SIZE)
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            shared_ira[ptr] = state->ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    ret = make_rng_int4(0,0,0,0);
    for(int nw=0; nw<4; ++nw) {
        rng_int_t rng = 0;
        for(i=0; i<bdim; i+=BB) {
            int off = tid - i;
            uchar4 ip = shared_ip;
            if(0 <= off && off < BB) {
                rng_int_t tmp = shared_ira[ip.y+off]+shared_ira[ip.z+off];
                shared_ira[ip.x+off] = tmp;
                rng = tmp ^ shared_ira[ip.w+off];
                // shift values
                ret.x = ret.y;
                ret.y = ret.z;
                ret.z = ret.w;
                ret.w = rng;
                //ip = make_uchar4(ip.x+off,ip.y+off,ip.z+off,ip.w+off);
            }
            __syncthreads();

            if(off == 0)
                shared_ip = make_uchar4(ip.x+BB,ip.y+BB,ip.z+BB,ip.w+BB);
            __syncthreads();
        }
    }

    // copy back shared state
    if(tid == 0) {
        state->ip = shared_ip;
    }
    __syncthreads();
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            state->ira[ptr] = shared_ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    return ret;
}

//-----------------------------------------------------------------------------

__device__ void rng_uniform_shm16_2xint4(rng_state* states, rng_int4_t &ret1, rng_int4_t &ret2)
{
    const int tid = threadIdx.x;
    const int bdim = blockDim.x;
    int i;
    rng_state* state = &states[blockIdx.x];
    rng_int4_t ret;
    __shared__ int shared_ira[RNG_N];
    __shared__ uchar4 shared_ip;
    int bs;

    // copy state to shared state
    if(tid == 0) {
        shared_ip = state->ip;
    }
    __syncthreads();

    // loop only if RNG_N > bdim(==BLOCK_SIZE)
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            shared_ira[ptr] = state->ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);

    for(int nb=0; nb<2; ++nb) {
        ret = make_rng_int4(0,0,0,0);
        for(int nw=0; nw<4; ++nw) {
            rng_int_t rng = 0;
            for(i=0; i<bdim; i+=BB) {
                int off = tid - i;
                uchar4 ip = shared_ip;
                if(0 <= off && off < BB) {
                    rng_int_t tmp = shared_ira[ip.y+off]+shared_ira[ip.z+off];
                    shared_ira[ip.x+off] = tmp;
                    rng = tmp ^ shared_ira[ip.w+off];
                    // shift values
                    ret.x = ret.y;
                    ret.y = ret.z;
                    ret.z = ret.w;
                    ret.w = rng;
                    //ip = make_uchar4(ip.x+off,ip.y+off,ip.z+off,ip.w+off);
                }
                __syncthreads();

                if(off == 0)
                    shared_ip = make_uchar4(ip.x+BB,ip.y+BB,ip.z+BB,ip.w+BB);
                __syncthreads();
            }
        }
        ret1 = ret2;
        ret2 = ret;
    }

    // copy back shared state
    if(tid == 0) {
        state->ip = shared_ip;
    }
    __syncthreads();
    bs = 0;
    do {
        int ptr = bs*bdim+tid;
        if(ptr < RNG_N) {
            state->ira[ptr] = shared_ira[ptr];
        }
        __syncthreads();
    } while(++bs < RNG_N/bdim);
}

//-----------------------------------------------------------------------------

__device__ float rng_uniform_float(rng_state* states)
{
    rng_int_t tmp;
    float ret;
    tmp = rng_uniform_int(states);
    //tmp = states->bu;
    ret = (float)tmp/(float)(RNG_INT_MAX);  
    return ret;
}

//-----------------------------------------------------------------------------

__device__ float rng_uniform_shm_float(rng_state* states)
{
    rng_int_t tmp;
    float ret;
    tmp = rng_uniform_shm_int(states);
    //tmp = states->bu;
    ret = (float)tmp/(float)(RNG_INT_MAX);  
    return ret;
}

//-----------------------------------------------------------------------------

__device__ float rng_uniform_shm16_float(rng_state* states)
{
    rng_int_t tmp;
    float ret;
    tmp = rng_uniform_shm16_int(states);
    //tmp = states->bu;
    ret = (float)tmp/(float)(RNG_INT_MAX);  
    return ret;
}

//-----------------------------------------------------------------------------

__device__ float2 rng_uniform_shm16_float2(rng_state* states)
{
    rng_int2_t tmp;
    tmp = rng_uniform_shm16_int2(states);
    //tmp = states->bu;
    return make_float2((float)tmp.x/(float)(RNG_INT_MAX), (float)tmp.y/(float)(RNG_INT_MAX));
}

//-----------------------------------------------------------------------------

__device__ float4 rng_uniform_shm16_float4(rng_state* states)
{
    rng_int4_t tmp;
    tmp = rng_uniform_shm16_int4(states);
    //tmp = states->bu;
    return make_float4((float)tmp.x/(float)(RNG_INT_MAX),
                       (float)tmp.y/(float)(RNG_INT_MAX),
                       (float)tmp.z/(float)(RNG_INT_MAX),
                       (float)tmp.w/(float)(RNG_INT_MAX));
}

//-----------------------------------------------------------------------------

__device__ void rng_uniform_shm16_2xfloat4(rng_state* states, float4 &ret1, float4 &ret2)
{
    rng_int4_t tmp1, tmp2;
    rng_uniform_shm16_2xint4(states, tmp1, tmp2);
    //tmp = states->bu;
    ret1 = make_float4((float)tmp1.x/(float)(RNG_INT_MAX),
                       (float)tmp1.y/(float)(RNG_INT_MAX),
                       (float)tmp1.z/(float)(RNG_INT_MAX),
                       (float)tmp1.w/(float)(RNG_INT_MAX));
    ret2 = make_float4((float)tmp2.x/(float)(RNG_INT_MAX),
                       (float)tmp2.y/(float)(RNG_INT_MAX),
                       (float)tmp2.z/(float)(RNG_INT_MAX),
                       (float)tmp2.w/(float)(RNG_INT_MAX));
}

//-----------------------------------------------------------------------------

__device__ void rng_uniform_shm16_4xfloat2(rng_state* states, float2 &rng0, float2 &rng1, float2 &rng2, float2 &rng3)
{
    rng_int4_t tmp1, tmp2;
    rng_uniform_shm16_2xint4(states, tmp1, tmp2);

    rng0 = make_float2((float)tmp1.x/(float)(RNG_INT_MAX),
                       (float)tmp1.y/(float)(RNG_INT_MAX));
    rng1 = make_float2((float)tmp1.z/(float)(RNG_INT_MAX),
                       (float)tmp1.w/(float)(RNG_INT_MAX));
    rng2 = make_float2((float)tmp2.x/(float)(RNG_INT_MAX),
                       (float)tmp2.y/(float)(RNG_INT_MAX));
    rng3 = make_float2((float)tmp2.z/(float)(RNG_INT_MAX),
                       (float)tmp2.w/(float)(RNG_INT_MAX));
}

//-----------------------------------------------------------------------------
// Gaussian
//-----------------------------------------------------------------------------

__device__ float rng_gauss_float(rng_state* states)
{
    float2 tmp;
    tmp = rng_uniform_shm16_float2(states);
    return sqrt(-2.0*(log(tmp.x))) * cos(2.0 * CUDART_PI_F * tmp.y);
}

//-----------------------------------------------------------------------------

__device__ float2 rng_gauss_float2(rng_state* states)
{
    float2 tmp;
    float t;
    tmp = rng_uniform_shm16_float2(states);
    t = sqrt(-2.0*(log(tmp.x)));
    return make_float2(t * cos(2.0 * CUDART_PI_F * tmp.y),
                       t * sin(2.0 * CUDART_PI_F * tmp.y));

}

//-----------------------------------------------------------------------------
//-----------------------------------------------------------------------------
/*
 * Local variables:
 *  mode: c++
 *  c-indent-level: 4
 *  c-basic-offset: 4
 *  tab-width: 4
 *  indent-tabs-mode: nil
 * End:
 */

#endif // __CUHB_RNG_CUH__
